LlamaIndexのインデックスをディスクに保存・再ロードしてみた

LlamIndexで作成したGPTVectorStoreIndexのインデックスをストレージに書き出すことで、コンピューティングリソースをまたいで共有してみました。

#LlamaIndex

#OpenAI

#Python

nayu.t.s

2023.06.26

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

データアナリティクス事業本部インテグレーション部機械学習チームの鈴木です。

ある環境で作成したLlamaIndexのインデックスを一旦ディスクに永続化し、別の環境で再ロードする手順を試してみたのでご紹介します。

本記事の内容

デフォルトではLlamaIndexはデータをメモリ内に保存するため、ディスクに保存しておくことで、簡単に保管・再読み込みできるようにしたいということが目的です。

以下のLlamaIndexの『Persisting & Loading Data』のドキュメントを参考にしました。

このドキュメントに記載の内容のうち、GPTVectorStoreIndexのpersistによるディスクへの保存と、load_indices_from_storageによる読み込みを試しました。

なお、『AttributeError: 'GPTVectorStoreIndex' object has no attribute 'save_to_disk' · Issue #4581 · jerryjliu/llama_index』によると、v0.6.0からこのAPIに変更されているようでした。

クライアント実行環境

実行環境はGoogle Colaboratoryを使いました。ハードウェアアクセラレータ無し、ランタイム仕様は標準としました。

Pythonのバージョンは以下でした。

!python --version
# Python 3.10.12

また、ライブラリは以下のようにインストールしました。

# 保存・読み込み両環境で使用
!pip install llama-index
!pip install python-dotenv

# 保存環境でSimpleWebPageReaderにて使用
!pip install html2text

インストールされたライブラリのバージョンは以下でした。

pip freeze | grep -e "openai" -e "llama" -e "langchain"
# langchain==0.0.214
# langchainplus-sdk==0.0.17
# llama-index==0.6.32
# openai==0.27.8

環境変数の設定

まず、必要な値を.envファイルに書き込みます。

!echo 'OPENAI_API_KEY="<トークン>"' >.env

load_dotenv()で環境変数を読み込みました。

from dotenv import load_dotenv
load_dotenv()

やってみる

1. インデックスの作成から保存まで

あるGoogle Colabのセッションで、先述の設定を行った後、以下の順でzipファイルの作成までを行いました。

まず、ライブラリを読み込みました。

from llama_index import GPTVectorStoreIndex
from llama_index import ServiceContext
from llama_index import SimpleWebPageReader
from llama_index import StorageContext
from llama_index.callbacks import CallbackManager, LlamaDebugHandler
from llama_index.storage.docstore import SimpleDocumentStore
from llama_index.storage.index_store import SimpleIndexStore
from llama_index.vector_stores import SimpleVectorStore

続いて、SimpleWebPageReaderで機械学習チームで提供しているサービスの紹介ページを読み込みました。

documents = SimpleWebPageReader(html_to_text=True).load_data([
    "https://classmethod.jp/services/machine-learning/",
    "https://classmethod.jp/services/machine-learning/data-assessment/",
    "https://classmethod.jp/services/machine-learning/recommend/"
])

GPTVectorStoreIndexのインスタンスを作成しました。今回はデバッグ用にCallbackManagerを設定しておきました。

# デバッグ用の設定
llama_debug_handler = LlamaDebugHandler()
callback_manager = CallbackManager([llama_debug_handler])
service_context = ServiceContext.from_defaults(callback_manager=callback_manager)

# GPTVectorStoreIndexの作成
vector_store_index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)

# **********
# Trace: index_construction
#     |_node_parsing ->  0.089474 seconds
#       |_chunking ->  0.031098 seconds
#       |_chunking ->  0.030921 seconds
#       |_chunking ->  0.026291 seconds
#     |_embedding ->  0.717604 seconds
#     |_embedding ->  0.509981 seconds
# **********

保存用のディレクトリを作成しました。

!mkdir ./storage_context

persistでstorage_contextを保存しました。

storage_context = vector_store_index.storage_context
storage_context.persist(persist_dir="./storage_context")

zipコマンドで圧縮しました。

!zip -r storage_context.zip "./storage_context"

以下のようにzipファイル化できました。

保存したストレージコンテキスト

zipファイルは手動でローカルPCにダウンロードしておきました。

ちなみに、ノードの分割状況は以下のようになっていました。

for doc_id, node in vector_store_index.storage_context.docstore.docs.items():
    node_dict = node.to_dict()
    print(f'{doc_id=}, len={len(node_dict["text"])}, start={node_dict["node_info"]["start"]}, end={node_dict["node_info"]["end"]}')

# doc_id='a5c79a2c-279a-47ad-b4d2-522921188427', len=1964, start=0, end=1964
# doc_id='ee4a9f12-b6d6-45b5-85b3-5af2ca3993d4', len=1123, start=1965, end=3088
# doc_id='a94e81cb-c8db-4999-8bd9-17e412ca2a94', len=860, start=3089, end=3949
# doc_id='d9b71f34-252d-425d-b1ec-c5b2bb45fe11', len=1345, start=3931, end=5276
# doc_id='06a52ff2-85b2-4392-a2b2-2484cae7ce74', len=1068, start=5296, end=6364
# doc_id='84ca182c-58b3-4c5f-b156-c6f7d527577a', len=2082, start=6365, end=8447
# doc_id='c367295f-e4b0-4169-bd59-8457488742af', len=1774, start=8438, end=10212
# doc_id='c9ee4b25-39ca-4721-91bf-0d4e747fefb1', len=1964, start=0, end=1964
# doc_id='b2eea9d4-a8a9-4179-a539-3196ae0cf4ed', len=1057, start=1965, end=3022
# doc_id='eafbcb3a-339a-49e3-aafa-d9c9a4e88e84', len=856, start=3023, end=3879
# doc_id='66732644-4bca-4b06-ad57-8a85ced37f31', len=1805, start=3880, end=5685
# doc_id='06bec6bf-2c4a-4059-86f3-407e15c0ae21', len=2235, start=5686, end=7921
# doc_id='b44f8442-9abc-4c6f-9819-5c0511b20d30', len=65, start=7922, end=7987
# doc_id='491d5eac-0cfd-41b2-b485-95353c3e72f2', len=1964, start=0, end=1964
# doc_id='cbf004c3-1219-457b-b479-6e57ee023a32', len=1165, start=1965, end=3130
# doc_id='47955b85-e3fb-4d03-b5a6-74cf8a3d6930', len=1020, start=3131, end=4151
# doc_id='1475c828-9e4a-425e-88c8-48a01b464eec', len=955, start=4139, end=5094
# doc_id='38c0db1c-e568-4e2b-8cfb-dceca403ae64', len=1588, start=5108, end=6696
# doc_id='c1634ed7-4866-407b-a096-a17f2cfcb23b', len=2234, start=6697, end=8931
# doc_id='99e3aa00-e5fe-4d96-a681-8a9322a86e30', len=409, start=8932, end=9341

2. インデックスの読み込みから回答生成まで

次に、手順1で使ったものとは異なるGoogle Colabのセッションで、先述の設定を行った後、以下の順で回答生成までを行いました。

まず、ライブラリを読み込みました。

from llama_index.callbacks import CallbackManager, LlamaDebugHandler
from llama_index import load_index_from_storage
from llama_index import ServiceContext
from llama_index.storage.docstore import SimpleDocumentStore
from llama_index.storage.index_store import SimpleIndexStore
from llama_index.vector_stores import SimpleVectorStore
from llama_index import StorageContext

ローカルPCからzipファイルをアップロードした後、unzipコマンドで解凍しました。

!unzip ./storage_context.zip

以下のようなディレクトリ構成になりました。

アップロードしたストレージコンテキスト

デバッグ用にCallbackManagerを設定したservice_contextインスタンスを作成しました。

llama_debug_handler = LlamaDebugHandler()
callback_manager = CallbackManager([llama_debug_handler])
service_context = ServiceContext.from_defaults(callback_manager=callback_manager)

ドキュメントを参考に、storage_contextインスタンスを作成しました。このとき、from_persist_dirでzipファイルを展開してできたディレクトリを指定しました。

load_index_from_storageでVectorStoreIndexのインスタンスを作成しました。

storage_context = StorageContext.from_defaults(
    docstore=SimpleDocumentStore.from_persist_dir(persist_dir="./storage_context"),
    vector_store=SimpleVectorStore.from_persist_dir(persist_dir="./storage_context"),
    index_store=SimpleIndexStore.from_persist_dir(persist_dir="./storage_context"),
)

# don't need to specify index_id if there's only one index in storage context
vector_store_index = load_index_from_storage(storage_context, service_context=service_context)

機械学習支援サービスの内容ついて教えてください。というクエリで回答を生成してみました。

query_engine = vector_store_index.as_query_engine(service_context=service_context)

response = query_engine.query("機械学習支援サービスの内容ついて教えてください。")

# **********
# **********
# Trace: query
#     |_query ->  12.342844 seconds
#       |_retrieve ->  0.142348 seconds
#         |_embedding ->  0.13444 seconds
#       |_synthesize ->  12.200365 seconds
#         |_llm ->  12.194351 seconds
# **********

以下のように読み込んだドキュメントをもとに回答が生成されたことを確認できました！

for i in response.response.split("。"):
    print(i + "。")

# クラスメソッドが提供する機械学習支援サービスは、Amazon SageMakerなどAWSの機械学習サービスを活用したシステムの導入支援を行います。
# ECサイトのレコメンドシステムやテキストマイニングを活用したインサイトの発見、改善に必要なアクションがわかるなど、様々なケースを支援できます。
# また、クラスメソッドが持つ機械学習の豊富な知見をもとに、課題に対する優先順位付けの実施と必要となるアクションを提示します。
# 。

クラスメソッドが提供する機械学習支援サービスは、Amazon SageMakerなどAWSの機械学習サービスを活用したシステムの導入支援を行います。

ECサイトのレコメンドシステムやテキストマイニングを活用したインサイトの発見、改善に必要なアクションがわかるなど、様々なケースを支援できます。

また、クラスメソッドが持つ機械学習の豊富な知見をもとに、課題に対する優先順位付けの実施と必要となるアクションを提示します。

ちなみに、ノードの分割状況は以下のようになっていました。これはzipファイルを作成した側のセッションで確認したものと同じですね。

for doc_id, node in vector_store_index.storage_context.docstore.docs.items():
    node_dict = node.to_dict()
    print(f'{doc_id=}, len={len(node_dict["text"])}, start={node_dict["node_info"]["start"]}, end={node_dict["node_info"]["end"]}')

# doc_id='a5c79a2c-279a-47ad-b4d2-522921188427', len=1964, start=0, end=1964
# doc_id='ee4a9f12-b6d6-45b5-85b3-5af2ca3993d4', len=1123, start=1965, end=3088
# doc_id='a94e81cb-c8db-4999-8bd9-17e412ca2a94', len=860, start=3089, end=3949
# doc_id='d9b71f34-252d-425d-b1ec-c5b2bb45fe11', len=1345, start=3931, end=5276
# doc_id='06a52ff2-85b2-4392-a2b2-2484cae7ce74', len=1068, start=5296, end=6364
# doc_id='84ca182c-58b3-4c5f-b156-c6f7d527577a', len=2082, start=6365, end=8447
# doc_id='c367295f-e4b0-4169-bd59-8457488742af', len=1774, start=8438, end=10212
# doc_id='c9ee4b25-39ca-4721-91bf-0d4e747fefb1', len=1964, start=0, end=1964
# doc_id='b2eea9d4-a8a9-4179-a539-3196ae0cf4ed', len=1057, start=1965, end=3022
# doc_id='eafbcb3a-339a-49e3-aafa-d9c9a4e88e84', len=856, start=3023, end=3879
# doc_id='66732644-4bca-4b06-ad57-8a85ced37f31', len=1805, start=3880, end=5685
# doc_id='06bec6bf-2c4a-4059-86f3-407e15c0ae21', len=2235, start=5686, end=7921
# doc_id='b44f8442-9abc-4c6f-9819-5c0511b20d30', len=65, start=7922, end=7987
# doc_id='491d5eac-0cfd-41b2-b485-95353c3e72f2', len=1964, start=0, end=1964
# doc_id='cbf004c3-1219-457b-b479-6e57ee023a32', len=1165, start=1965, end=3130
# doc_id='47955b85-e3fb-4d03-b5a6-74cf8a3d6930', len=1020, start=3131, end=4151
# doc_id='1475c828-9e4a-425e-88c8-48a01b464eec', len=955, start=4139, end=5094
# doc_id='38c0db1c-e568-4e2b-8cfb-dceca403ae64', len=1588, start=5108, end=6696
# doc_id='c1634ed7-4866-407b-a096-a17f2cfcb23b', len=2234, start=6697, end=8931
# doc_id='99e3aa00-e5fe-4d96-a681-8a9322a86e30', len=409, start=8932, end=9341